import argparse
import json
import logging
import os
import pathlib
import subprocess
import sys
from logging.handlers import TimedRotatingFileHandler


def init_log(name='dlutfem', config: dict = None):
    if config is None:
        config = {}

    _logger = logging.getLogger(name)
    if len(_logger.handlers) > 0:
        return _logger

    log_file = config.get('filePath')
    level = config.get('level', 'INFO')
    backup_count = config.get('maxBackups', '7')

    _logger.setLevel(level)
    format_str = logging.Formatter('%(asctime)-15s %(name)s %(process)d %(levelname)-8s %(message)s')

    stream_handler = logging.StreamHandler(stream=sys.stdout)
    stream_handler.setFormatter(format_str)
    _logger.addHandler(stream_handler)

    if log_file:
        log_path = pathlib.Path(log_file).parent
        log_path.mkdir(parents=True, exist_ok=True)
        rotate_handler = TimedRotatingFileHandler(log_file, when='D', backupCount=int(backup_count))
        rotate_handler.setFormatter(format_str)
        _logger.addHandler(rotate_handler)

    return _logger


def load_driver_config(config_path):
    cp = pathlib.Path(config_path)
    if not cp.is_absolute():
        cp = pathlib.Path(os.getcwd(), config_path)

    if not cp.is_file():
        raise Exception('driver.config "{}" not exists'.format(config_path))

    with open(cp, 'r') as f:
        config = json.load(f)
        if config.get('product') != 'dlutfem':
            raise Exception('invalid driver.config')
        return config


def run_simulation(args, driver_cfg):
    with open(args.sim_config, 'r', encoding="utf-8") as file:
        sim_config = json.load(file)
    outputFilePath = args.output_file_path
    if outputFilePath:
        sim_config['outputFilePath'] = outputFilePath
    logger.info(sim_config)
    dlutfem_image = driver_cfg.get("dlutfemImage")
    # model_entry_path = os.path.join(sim_config['model']['file'], sim_config['model']['entry'])
    model_path = sim_config['model']['file']
    model_file = sim_config['model']['entry']
    result_path = pathlib.Path(os.path.join(sim_config['model']['file'], 'results'))
    result_path.mkdir(parents=True, exist_ok=True)
    # docker run
    # 构建Docker命令
    cmd = [
        'docker', 'run', '--rm', '--privileged',
        '-v', f'{model_path}:/home/models',
        dlutfem_image, 'SGFEM', '-i', f'/home/models/{model_file}'
    ]

    logger.info(f'Starting simulation with command: {" ".join(cmd)}')

    try:
        with subprocess.Popen(
                cmd,
                shell=False,  # 更安全的做法是不使用shell
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                encoding='utf-8',  # 替代universal_newlines
                cwd=str(sim_config['model']['file']),  # 确保路径是字符串
                errors='replace'  # 处理编码错误
        ) as p:
            # 实时读取输出，避免缓冲区满
            while True:
                output = p.stdout.readline()
                error = p.stderr.readline()

                if output == '' and error == '' and p.poll() is not None:
                    break
                if output:
                    logger.info(f'STDOUT: {output.strip()}')
                if error:
                    logger.error(f'STDERR: {error.strip()}')
            return_code = p.poll()
            if return_code != 0:
                logger.error(f'Simulation failed with return code: {return_code}')
                raise subprocess.CalledProcessError(return_code, cmd)
            logger.info('Simulation completed successfully')
    except FileNotFoundError as e:
        logger.error(f'Executable not found: {e}')
        raise
    except subprocess.SubprocessError as e:
        logger.error(f'Subprocess error: {e}')
        raise
    except Exception as e:
        logger.error(f'Unexpected error during simulation: {e}')
        raise
    output = {
        "success": True,
        "message": "",
        "data": {
            "simulationResultDir": model_path,
            "specifiedParameters": []
        }
    }
    # 驱动运行结果写入output.json文件
    with open(outputFilePath, 'w') as f:
        json.dump(output, f)


def main(args):
    global logger
    config = load_driver_config(args.driver_config)
    driver_cfg = config.get('config', {})
    logger = init_log(config=driver_cfg.get('log', {}))
    logger.info(args)

    # 运行模型
    if args.subcommand == "run":
        return run_simulation(args, driver_cfg)
    elif args.subcommand == "parse":
        raise Exception('parse not support')
    else:
        raise Exception('invalid sub command')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='driver for dlutfem')

    subparsers = parser.add_subparsers(title='Available Commands', dest="subcommand", help='sub-command is required')
    parser_run = subparsers.add_parser('run', help='run simulation')
    parser_run.add_argument('--driver-config', '-d', type=str, default='driver.config.json', help='driver config')
    parser_run.add_argument('--output-file-path', '-o', type=str, required=True, help='output file path')
    parser_run.add_argument('--sim-config', '-c', type=str, required=True, help='config for simulation task')

    parser_parse = subparsers.add_parser('parse', help='parse configurable parameters')
    parser_parse.add_argument('--output-file-path', '-o', type=str, required=True, help='output file path')
    parser_parse.add_argument('--driver-config', '-d', type=str, default='driver.config.json', help='driver config')
    parser_parse.add_argument('--model-dir', '-m', type=str, required=True, help='model dir')
    parser_parse.add_argument('--entry', '-e', type=str, required=True, help='model file')

    args_ = parser.parse_args()
    if not args_.subcommand:
        parser.print_help()
        sys.exit(1)
    sys.exit(main(args_))